#
# 模型计算的库
# 
import cython
cimport cython # 必须导入
import numpy as np
cimport numpy as np
from libc.math cimport pi
from scipy.optimize import leastsq
import random
from tqdm import tqdm # 进度条
import logging
logger = logging.getLogger("mylog")


def WMCModel(param_arr,sample_lai,sample_soil,sample_inc,sample_sigma,sample_soil_dbl):
    """ WMC模型

    Args:
        param_arr (np.ndarray): 参数数组
        sample_lai (double): 叶面积指数
        sample_soil (double): 土壤含水量
        sample_inc (double): 入射角（弧度值）
        sample_sigma (double): 后向散射系数（线性值）

    Returns:
        double: 方程值
    """
    # 映射参数，方便修改模型
    A,B,C,D,M,N=param_arr  # 在这里修改模型
    V_lai=sample_lai
    #V_lai=E*sample_lai+F
    exp_gamma=np.exp(-2*B*((V_lai*D+C))*(1/np.cos(sample_inc)))
    sigma_soil=sample_soil_dbl*0+(M*sample_soil+N)
    sigma_veg=A*((V_lai))*np.cos(sample_inc)
    result=sigma_veg*(1-exp_gamma)+sigma_soil*exp_gamma-sample_sigma
    return result

def train_WMCmodel(lai_waiter_inc_sigma_list,params_X0,train_err_image_path,draw_flag=True):
    """ 训练模型参数

    Args:
        lai_waiter_inc_sigma_list (list): 训练模型使用的样本呢
    """
    def f(X):
        eqs=[]
        for lai_waiter_inc_sigma_item in lai_waiter_inc_sigma_list:
            #logger.info(lai_waiter_inc_sigma_item)
            sample_time,sample_code,sample_lon,sample_lat,sample_lai,csv_sigma,sample_soil,sample_inc,sample_sigma,sample_soil_dbl=lai_waiter_inc_sigma_item
            #logger.info(str(sample_lai,csv_sigma))
            eqs.append(WMCModel(X,sample_lai,sample_soil,sample_inc,csv_sigma,sample_soil_dbl))
        return eqs
    
    X0 = params_X0 # 初始值
    logger.info(str(X0))
    h = leastsq(f, X0)
    logger.info(h[0],h[1])
    err_f=f(h[0])
    x_arr=[lai_waiter_inc_sigma_item[4] for lai_waiter_inc_sigma_item in lai_waiter_inc_sigma_list]
    # 根据误差大小进行排序
    logger.info("训练集：\n根据误差输出点序\n数量：{}\n点序\t误差值\t 样点信息".format(str(np.array(err_f).shape)))
    for i in np.argsort(np.array(err_f)):
        logger.info('{}\t{}\t{}'.format(i,err_f[i],str(lai_waiter_inc_sigma_list[i])))
    logger.info("\n误差点序输出结束\n")

    if draw_flag:
        logger.info(err_f)
        logger.info(np.where(np.abs(err_f)<10))
        from matplotlib import pyplot as plt
        plt.scatter(x_arr,err_f)
        plt.title("equation-err")
        plt.savefig(train_err_image_path,dpi=600)
        plt.show()
        
    return h[0]

def test_WMCModel(lai_waiter_inc_sigma_list,param_arr,lai_X0,test_err_image_path,draw_flag=True):
    """ 测试模型训练结果

    Args:
        lai_waiter_inc_sigma_list (list): 测试使用的样本集
        A (_type_): 参数A
        B (_type_): 参数B
        C (_type_): 参数C
        D (_type_): 参数D
        M (_type_): 参数M
        N (_type_): 参数N
        lai_X0 (_type_): 初始值

    Returns:
        list: 误差列表 [sample_lai,err,predict]
    """
    err=[]
    err_f=[]
    x_arr=[]
    err_lai=[]
    for lai_waiter_inc_sigma_item in lai_waiter_inc_sigma_list:
        sample_time,sample_code,sample_lon,sample_lat,sample_lai,csv_sigma,sample_soil,sample_inc,sample_sigma,sample_soil_dbl=lai_waiter_inc_sigma_item
        def f(X):
            lai=X[0]
            eqs=[WMCModel(param_arr,lai,sample_soil,sample_inc,csv_sigma,sample_soil_dbl)]
            return eqs
        X0=lai_X0
        h = leastsq(f, X0)
        temp_err=h[0]-sample_lai
        err_lai.append(temp_err[0]) # lai预测的插值
        err.append([sample_lai,temp_err[0],h[0][0],sample_code])
        err_f.append(f(h[0])[0]) # 方程差
        x_arr.append(sample_lai)
    
    # 根据误差大小进行排序
    logger.info("测试集：\n根据误差输出点序\n数量：{}\n点序\t误差值\t 方程差\t样点信息".format(str(np.array(err_lai).shape)))
    for i in np.argsort(np.array(err_lai)):
        logger.info('{}\t{}\t{}\t{}'.format(i,err_lai[i],err_f[i],str(lai_waiter_inc_sigma_list[i])))
    logger.info("\n误差点序输出结束\n")

    if draw_flag:
        from matplotlib import pyplot as plt
        plt.scatter(x_arr,err_lai)
        plt.title("equation-err")
        plt.savefig(test_err_image_path,dpi=600)
        plt.show() 
    return err

def processs_WMCModel(param_arr,lai_X0,sigma,inc_angle,soil_water,sample_soil_dbl):
    if(sigma<0 or sample_soil_dbl <0):
        return np.nan
    def f(X):
        lai=X[0]
        eqs=[WMCModel(param_arr,lai,soil_water,inc_angle,sigma,sample_soil_dbl)]
        return eqs
    h = leastsq(f, [lai_X0])

    return h[0][0]

# Cython 的扩展地址
cpdef np.ndarray[double,ndim=2] process_tiff(np.ndarray[double,ndim=2] sigma_tiff,
                                             np.ndarray[double,ndim=2] inc_tiff,
                                             np.ndarray[double,ndim=2] soil_water_tiff,
                                             np.ndarray[double,ndim=2] sample_soil_dbl_tiff,
                                             np.ndarray[double,ndim=1] param_arr,
                                             double lai_X0):
    
    cdef np.ndarray[double,ndim=2] result=sigma_tiff
    cdef int param_arr_length=param_arr.shape[0]
    cdef int height=sigma_tiff.shape[0]
    cdef int width=sigma_tiff.shape[1]
    cdef int i=0
    cdef int j=0
    cdef double temp=0
    pbar = tqdm(total=height) # 添加进度条
    while i<height:
        j=0
        while j<width:
            temp = processs_WMCModel(param_arr,lai_X0,sigma_tiff[i,j],inc_tiff[i,j],soil_water_tiff[i,j],sample_soil_dbl_tiff[i,j])
            temp=temp if temp<10 and temp>=0 else np.nan
            result[i,j]=temp
            j=j+1
        i=i+1
        pbar.update(1) # 更新进度
    pbar.close()
    return result
    